'''
Created on Nov 15, 2013

@author: maxp
'''
import sys
import argparse
from Bio import SeqIO as bseq
from numpy import arange, zeros
import align
#from progressbar import ETA, ProgressBar, SimpleProgress, Percentage, Bar, FileTransferSpeed

#===================================================================================================
# LCP(s1, s2)
# -computes the length of the longest common prefix between two strings s1 and s2
# -helper function to reindex the DP table
#===================================================================================================
def LCP(str1, str2):
    for i in range(0,min([len(str1),len(str2)])):
        if str1[i] != str2[i]:
            return i
    return i + 1

def main():
    
    
    #===============================================================================================
    # This section is just argument parsing from command line
    #===============================================================================================
    
    parser = argparse.ArgumentParser(description="For each read in 'reference' find all reads in 'query' within an edit distance")
    parser.add_argument('ref_file', nargs='?', default="reference.fna",
                        help='File containing the shorter list of reads.')
    parser.add_argument('query_file', nargs='?', default="query.fna",
                        help='File containing longer list of query reads.')
    parser.add_argument('--out_file', nargs='?', const="ref_matches.txt", type=argparse.FileType('w'),
                        default=sys.stdout,
                        help='The destination of the output file (default is stdout).')
    parser.add_argument('--maxedit', default=6, type=int,
                        help='Maxmimum edit distance to consider.')
    
    args = parser.parse_args()
        

    #===============================================================================================
    # Here we initialize the DP table...
    # Since the longest query and reference in our dataset is 635 characters, this is the hardcoded
    # size of our table (i.e. longer sequences will result in error)
    #===============================================================================================
    T = zeros(shape=(635,635),dtype=int)
    T[0] = arange(635)
    T[:,0] = arange(635)

    #===============================================================================================
    # This line creates a list of (sequence, id) tuples for all sequences in the query file.
    # Python's timsort is used to alphabetically sort the list thus emulating the suffix tree
    #===============================================================================================
    qs_and_ids = sorted((str(rec.seq),rec.id) for rec in bseq.parse(args.query_file,"fasta"))
        
    #===============================================================================================
    # #===============================================================================================
    # # Pretty looking progress bar to track progress for each reference
    # #===============================================================================================
    # widgets = ['Progress: ', SimpleProgress(), '-', Percentage(), ' ',
    #            Bar(marker='0',left='[',right=']'), ETA(), FileTransferSpeed()]
    #===============================================================================================
                                                                      
    #===============================================================================================
    # Begin main loop (i.e. the short list of reference reads)
    #===============================================================================================
    for ref_rec in bseq.parse(args.ref_file,"fasta"):
        ref_seq = str(ref_rec.seq)
        result = [ref_rec.id]
        newindex = 1 #start on row 1 of DP table (not 0)
        qnum = 0 
        #===========================================================================================
        # start the progress bar
        #===========================================================================================
        #===========================================================================================
        # pbar = ProgressBar(widgets=widgets, maxval=len(qs_and_ids)).start()
        #===========================================================================================
        qseq, qid = qs_and_ids[qnum] #load first query sequece

        #===========================================================================================
        # begin the long, inner loop containing all queries
        #===========================================================================================
        while qnum < len(qs_and_ids):
            
            
            #=======================================================================================
            # We perform the DP algorith on the 2 sequences, but we shorten the longer sequence 
            # to the length of the shorter sequence + maxedit dist where possible.
            # Calling align() from these cases allow us to retain the original sequence with 
            # minimal additional processing
            # Align() returns both whether there was an alignment and the furthest row calculated
            # in the DP table
            #=======================================================================================
            if len(qseq) > len(ref_seq) + args.maxedit: #truncate query
                aligned, newindex = align.align(qseq[:len(ref_seq)+ args.maxedit],
                                                ref_seq, 
                                                T, newindex, args.maxedit)
            elif len(ref_seq) > len(qseq) + args.maxedit: #truncate reference
                aligned, newindex = align.align(qseq,
                                                ref_seq[:len(qseq)+args.maxedit],
                                                T, newindex, args.maxedit)
            else: #truncate neither
                aligned, newindex = align.align(qseq,
                                                ref_seq, 
                                                T, newindex, args.maxedit)
            
            #=======================================================================================
            # If the last query matched, then all querys with LCP > index will match
            # the LCP of the first query that doesn't match (+1) will be the new index
            #=======================================================================================
            if aligned: 
                result.append(qid)
            
            qnum += 1 #advance in the list of queries
            
            #=======================================================================================
            # If we exhausted all queries for this reference, we are done and can output the 
            # result for this reference to file
            # Otherwise, we need to advance through the query list to find the next candidate 
            # which needs algining...
            #=======================================================================================
            if qnum < len(qs_and_ids):
                new_qseq, new_qid = qs_and_ids[qnum] #look at next query
                
                #===================================================================================
                # This is the main timesaver of the algorithm...
                # If the LCP of the previous and the newsequence is greater than the deepest row
                # computed in the DP table, then whatever result we found with align() on the 
                # previous sequence, also holds for this new sequence.
                #===================================================================================
                while (LCP(qseq,new_qseq) > newindex):
                    if aligned: #if the old one aligned, then this one does too
                        result.append(new_qid)
                    qnum += 1 #increase the query index, but bound check before getting new query
                    if qnum < len(qs_and_ids):
                        new_qseq, new_qid = qs_and_ids[qnum]
                        #===========================================================================
                        # pbar.update(qnum)
                        #===========================================================================
                    else:
                        #===========================================================================
                        # Need to manually exit loop because loop guard is related to LCP which
                        # wont't change once we get to the last query  
                        #===========================================================================
                        break
                
                newindex = LCP(qseq,new_qseq) + 1
                qseq, qid = new_qseq, new_qid
        
        #===========================================================================================
        # Finish the progress bar, write the file and reset the result list for next reference
        #===========================================================================================
        #===========================================================================================
        # pbar.finish()        
        #===========================================================================================
        args.out_file.write('\t'.join(result))
        args.out_file.write('\n')
        print "finished with reference: ", result[0]
        del result
                  
    #===============================================================================================
    # All references finished.  Close the output file.
    #===============================================================================================
    args.out_file.close()        




if __name__ == '__main__':
    main()